#pragma once

#include <memory>
#include <string>

#include <storages/redis/impl/request.hpp>
#include <userver/storages/redis/base.hpp>
#include <userver/utils/assert.hpp>

#include <userver/storages/redis/client.hpp>
#include <userver/storages/redis/parse_reply.hpp>
#include <userver/storages/redis/request_data_base.hpp>

#include "client_impl.hpp"
#include "scan_reply.hpp"

USERVER_NAMESPACE_BEGIN

namespace storages::redis {

namespace impl {

void Wait(impl::Request& request);

template <ScanTag TScanTag>
storages::redis::Request<ScanReplyTmpl<TScanTag>> MakeScanRequest(
    ClientImpl& client,
    std::string key,
    size_t shard,
    typename ScanReplyTmpl<TScanTag>::Cursor cursor,
    ScanOptionsGeneric options,
    const CommandControl& command_control
) {
    if constexpr (TScanTag == ScanTag::kScan) {
        return client.MakeScanRequestNoKey(shard, cursor, std::move(options), command_control);
    } else {
        return client
            .MakeScanRequestWithKey<TScanTag>(std::move(key), shard, cursor, std::move(options), command_control);
    }
}

}  // namespace impl

template <typename Result, typename ReplyType>
class RequestDataImpl final : public RequestDataBase<ReplyType> {
public:
    explicit RequestDataImpl(impl::Request&& request)
        : request_(std::move(request))
    {}

    void Wait() override { impl::Wait(request_); }

    ReplyType Get(const std::string& request_description) override {
        auto reply = GetReply();
        return impl::ParseReply<Result, ReplyType>(std::move(reply), request_description);
    }

    ReplyPtr GetRaw() override { return GetReply(); }

    engine::impl::ContextAccessor* TryGetContextAccessor() noexcept override {
        return request_.TryGetContextAccessor();
    }

private:
    ReplyPtr GetReply() { return request_.Get(); }

    impl::Request request_;
};

template <typename Result, typename ReplyType>
class AggregateRequestDataImpl final : public RequestDataBase<ReplyType> {
    using RequestDataPtr = std::unique_ptr<RequestDataBase<ReplyType>>;

public:
    explicit AggregateRequestDataImpl(std::vector<RequestDataPtr>&& requests)
        : requests_(std::move(requests))
    {}

    void Wait() override {
        for (auto& request : requests_) {
            request->Wait();
        }
    }

    ReplyType Get(const std::string& request_description) override {
        std::vector<typename ReplyType::value_type> result;
        for (auto& request : requests_) {
            auto data = request->Get(request_description);
            std::move(data.begin(), data.end(), std::back_inserter(result));
        }
        return result;
    }

    ReplyPtr GetRaw() override {
        UASSERT_MSG(false, "Unsupported");
        return {};
    }

    engine::impl::ContextAccessor* TryGetContextAccessor() noexcept override {
        UASSERT_MSG(false, "Not implemented");
        return nullptr;
    }

private:
    std::vector<RequestDataPtr> requests_;
};

template <typename Result, typename ReplyType>
class DummyRequestDataImpl final : public RequestDataBase<ReplyType> {
public:
    explicit DummyRequestDataImpl(ReplyPtr&& reply)
        : reply_(std::move(reply))
    {}

    void Wait() override {}

    ReplyType Get(const std::string& request_description) override {
        return impl::ParseReply<Result, ReplyType>(std::move(reply_), request_description);
    }

    ReplyPtr GetRaw() override { return std::move(reply_); }

    engine::impl::ContextAccessor* TryGetContextAccessor() noexcept override {
        UASSERT_MSG(false, "Not implemented");
        return nullptr;
    }

private:
    ReplyPtr reply_;
};

template <ScanTag TScanTag>
class RequestScanData final : public RequestScanDataBase<TScanTag> {
public:
    using ReplyElem = typename ScanReplyElem<TScanTag>::type;

    template <ScanTag ScanTagParam = TScanTag>
    RequestScanData(
        std::shared_ptr<ClientImpl> client,
        size_t shard,
        ScanOptionsGeneric options,
        const CommandControl& command_control,
        std::enable_if_t<ScanTagParam == ScanTag::kScan>* = nullptr
    )
        : RequestScanData(std::move(client), {}, shard, std::move(options), command_control, TScanTag)
    {}

    template <ScanTag ScanTagParam = TScanTag>
    RequestScanData(
        std::shared_ptr<ClientImpl> client,
        std::string key,
        size_t shard,
        ScanOptionsGeneric options,
        const CommandControl& command_control,
        std::enable_if_t<ScanTagParam != ScanTag::kScan>* = nullptr
    )
        : RequestScanData(std::move(client), std::move(key), shard, std::move(options), command_control, TScanTag)
    {}

    ReplyElem Get() override;

    ReplyElem& Current() override;

    bool Eof() override;

private:
    RequestScanData(
        std::shared_ptr<ClientImpl> client,
        std::string key,
        size_t shard,
        ScanOptionsGeneric options,
        const CommandControl& command_control,
        ScanTag
    );

    void CheckReply();

    std::shared_ptr<ClientImpl> client_;
    std::string key_;
    size_t shard_{-1UL};
    ScanOptionsGeneric options_;
    CommandControl command_control_;

    using ScanReply = ScanReplyTmpl<TScanTag>;

    std::unique_ptr<Request<ScanReply>> request_;
    std::unique_ptr<ScanReply> reply_;
    size_t reply_keys_index_{0};
    bool eof_{false};
};

template <ScanTag TScanTag>
RequestScanData<TScanTag>::RequestScanData(
    std::shared_ptr<ClientImpl> client,
    std::string key,
    size_t shard,
    ScanOptionsGeneric options,
    const CommandControl& command_control,
    ScanTag
)
    : client_(std::move(client)),
      key_(std::move(key)),
      shard_(shard),
      options_(std::move(options)),
      command_control_(command_control),
      request_(std::make_unique<Request<
                   ScanReply>>(impl::MakeScanRequest<TScanTag>(*client_, key_, shard_, {}, options_, command_control_)))
{}

template <ScanTag TScanTag>
typename RequestScanData<TScanTag>::ReplyElem RequestScanData<TScanTag>::Get() {
    if (Eof()) {
        throw RequestScan::GetAfterEofException("Trying to Get() after eof");
    }
    UASSERT(reply_);
    UASSERT(reply_keys_index_ < reply_->GetKeys().size());
    return std::move(reply_->GetKeys()[reply_keys_index_++]);
}

template <ScanTag TScanTag>
typename RequestScanData<TScanTag>::ReplyElem& RequestScanData<TScanTag>::Current() {
    if (Eof()) {
        throw RequestScan::GetAfterEofException("Trying to call Current() after eof");
    }
    UASSERT(reply_);
    UASSERT(reply_keys_index_ < reply_->GetKeys().size());
    return reply_->GetKeys()[reply_keys_index_];
}

template <ScanTag TScanTag>
bool RequestScanData<TScanTag>::Eof() {
    CheckReply();
    return eof_;
}

template <ScanTag TScanTag>
void RequestScanData<TScanTag>::CheckReply() {
    while (!eof_ && (!reply_ || reply_keys_index_ == reply_->GetKeys().size())) {
        if (request_) {
            auto scan_reply_raw = request_->GetRaw();
            command_control_.force_server_id = scan_reply_raw->server_id;
            auto scan_reply = impl::ParseReply<ScanReply>(std::move(scan_reply_raw), this->request_description_);
            if (reply_) {
                *reply_ = std::move(scan_reply);
            } else {
                reply_ = std::make_unique<ScanReply>(std::move(scan_reply));
            }
        } else {
            reply_.reset();
            eof_ = true;
        }
        reply_keys_index_ = 0;

        if (!eof_ && reply_->GetCursor().GetValue()) {
            *request_ = impl::MakeScanRequest<
                TScanTag>(*client_, key_, shard_, reply_->GetCursor(), options_, command_control_);
        } else {
            request_.reset();
        }
    }
}

}  // namespace storages::redis

USERVER_NAMESPACE_END
